package com.test.tool.demo.websocket;

import cn.hutool.core.util.StrUtil;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;

import java.security.Principal;
import java.util.List;
import java.util.Map;

/**
 * 自定义消息发送拦截器
 *
 * @author 小八
 * @date 2022年04月25日15:43
 */
@Slf4j
public class WebSocketUserInterceptor implements ChannelInterceptor {

    /**
     * 登入统计处理器
     */
    private final OnlineStatisticalHandler handler;

    public WebSocketUserInterceptor(OnlineStatisticalHandler handler) {
        this.handler = handler;
    }

    @Override
    public void afterSendCompletion(Message<?> message, MessageChannel channel, boolean sent, Exception ex) {
        ChannelInterceptor.super.afterSendCompletion(message, channel, sent, ex);
    }

    /**
     * 获取包含在stomp中的用户信息
     */
    @SneakyThrows
    @SuppressWarnings("rawtypes")
    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
        Principal user = accessor.getUser();
        if (user == null) {
            Map<String, Object> attributes = (Map<String, Object>) accessor.getHeader("simpSessionAttributes");
            user = (Principal) attributes.get("user");
        }
        //1、判断是否首次连接
        if (accessor != null) {
            if(StompCommand.CONNECT.equals(accessor.getCommand())) {
                //2、判断token  账号类型:用户id   例如  agent:12345678910
                List<String> nativeHeader = accessor.getNativeHeader("userId");
                if (nativeHeader != null && !nativeHeader.isEmpty()) {
                    String token = nativeHeader.get(0);
                    if (StrUtil.isNotBlank(token)) {
                        //todo,通过token获取用户信息，下方用loginUser来代替
                        user = new WebSocketUser(token);
                        accessor.setUser(user);
                        Map<String, Object> attributes = (Map<String, Object>) accessor.getHeader("simpSessionAttributes");
                        attributes.put("user",user);
                        this.handler.statisticalOnline(user, OnlineStatisticalHandler.LOGIN_STATE);
                        return message;
                    }
                }
                throw new Exception("认证失败!");
            }
            if (StompCommand.DISCONNECT.equals(accessor.getCommand())) {
                this.handler.statisticalOnline(user, OnlineStatisticalHandler.OFFLINE_STATE);
            }
        }
        return message;
    }

}
